from selenium import webdriver
import time
from pymysql import connect


class DouyuSpider:

    def __init__(self):
        self.conn = connect(host='localhost', port=3306, user='root', password='root',
                            database='douyu_spider', charset='utf8')
        self.cursor = self.conn.cursor()
        self.start_url = 'https://www.douyu.com/directory/all'
        self.driver = webdriver.Chrome()

    # 解析数据
    def get_content_list(self):
        li_list = self.driver.find_elements_by_xpath('//ul[@id="live-list-contentbox"]/li')
        content_list = []
        for li in li_list:
            item = dict()
            item['room_img'] = li.find_element_by_xpath('.//span[@class="imgbox"]/img').get_attribute('src')
            item['room_name'] = li.find_element_by_xpath('./a').get_attribute('title')
            item['room_url'] = li.find_element_by_xpath('./a').get_attribute('href')
            item['room_user'] = li.find_element_by_xpath('.//span[@class="dy-name ellipsis fl"]').text
            item['room_cate'] = li.find_element_by_xpath('.//span[@class="tag ellipsis"]').text
            item['room_watches'] = li.find_element_by_xpath('.//span[@class="dy-num fr"]').text
            # print(item)
            content_list.append(item)
        # 获取下一页的url
        next_url = self.driver.find_elements_by_xpath('//a[@class="shark-pager-next"]')
        next_url = next_url[0] if len(next_url) > 0 else None
        return content_list, next_url

    # 保存分类，返回分类id
    def save_category(self, cate_name):
        # 保存分类
        cate_sql = 'select * from douyu_category where name = %s'
        params = [cate_name]
        self.cursor.execute(cate_sql, params)
        res = self.cursor.fetchone()
        if res is not None:
            # 如果已存在该分类，直接返回该分类id
            return res[0]
        else:
            # 如果不存在，就添加该分类
            sql = 'insert into douyu_category (name) values (%s)'
            self.cursor.execute(sql, params)
            self.conn.commit()
            return self.cursor.lastrowid

    # 观看人数转成数字
    def get_watchs_num(self, watches_str):
        if watches_str.find('万') > 0:
            return float(watches_str[:len(watches_str)-1]) * 10000
        else:
            return int(watches_str)

    # 保存房间数据
    def save_room(self, content_list):
        sql = "insert into douyu_room (room_cate, room_url, room_user, room_img, room_watches) values "
        for data in content_list:
            cate_id = self.save_category(data['room_cate'])
            watches_num = self.get_watchs_num(watches_str=data['room_watches'])
            params_sql = f"({cate_id},'{data['room_url']}','{data['room_user']}'," \
                         f"'{data['room_img']}',{watches_num}) ,"
            sql += params_sql
        sql = sql[:len(sql)-1]
        self.cursor.execute(sql)
        self.conn.commit()

    def run(self):
        # 发送请求，获取响应
        self.driver.get(self.start_url)
        # 提取数据，提取下一页的元素
        content_list, next_url = self.get_content_list()
        self.save_room(content_list)
        while next_url is not None:
            next_url.click()
            time.sleep(3)
            content_list, next_url = self.get_content_list()
            self.save_room(content_list)


if __name__ == '__main__':
    douyu = DouyuSpider()
    douyu.run()
